#!/usr/bin/env python3
"""
数据库初始化脚本
Database Initialization Script

用于初始化Neo4j数据库，导入医疗数据和知识图谱
Used to initialize Neo4j database with medical data and knowledge graph
"""

import os
import sys
import logging
import argparse
from datetime import datetime

# 添加项目根目录到系统路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

from api.services.neo4j_service import neo4j_service
from api.config.database import Config

def setup_logging():
    """设置日志"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler('database_init.log', encoding='utf-8')
        ]
    )

def check_neo4j_connection():
    """检查Neo4j连接"""
    try:
        status = neo4j_service.check_database_status()
        if 'error' in status:
            print(f"❌ Neo4j连接失败: {status['error']}")
            return False
        else:
            print("✅ Neo4j连接成功")
            return True
    except Exception as e:
        print(f"❌ Neo4j连接测试失败: {str(e)}")
        return False

def initialize_database(force=False):
    """初始化数据库"""
    print("="*60)
    print("🚀 开始初始化Neo4j数据库")
    print("="*60)
    
    # 检查连接
    if not check_neo4j_connection():
        print("请确保Neo4j服务正在运行并且配置正确")
        print(f"Neo4j URI: {Config.NEO4J_URI}")
        print(f"用户名: {Config.NEO4J_USERNAME}")
        return False
    
    # 检查是否已初始化
    if not force:
        if neo4j_service.is_initialized():
            print("⚠️  数据库已经初始化")
            choice = input("是否要重新初始化？这将清空所有现有数据 (y/N): ")
            if choice.lower() not in ['y', 'yes', '是']:
                print("初始化取消")
                return True
    
    try:
        # 获取数据文件路径（相对于项目根目录）
        project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
        data_dir = os.path.join(project_root, 'api', 'data')
        medical_data_path = os.path.join(data_dir, 'medical_data.json')
        knowledge_graph_path = os.path.join(data_dir, 'knowledge_graph.json')
        
        # 检查文件是否存在
        if not os.path.exists(medical_data_path):
            print(f"❌ 医疗数据文件不存在: {medical_data_path}")
            return False
        
        if not os.path.exists(knowledge_graph_path):
            print(f"❌ 知识图谱文件不存在: {knowledge_graph_path}")
            return False
        
        print(f"📁 医疗数据文件: {medical_data_path}")
        print(f"📁 知识图谱文件: {knowledge_graph_path}")
        
        # 开始初始化
        start_time = datetime.now()
        print(f"⏰ 开始时间: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")
        
        neo4j_service.initialize_database(
            medical_data_path=medical_data_path,
            knowledge_graph_path=knowledge_graph_path
        )
        
        end_time = datetime.now()
        duration = end_time - start_time
        
        print(f"⏰ 完成时间: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
        print(f"⏱️  耗时: {duration.total_seconds():.2f}秒")
        
        # 显示初始化结果
        status = neo4j_service.check_database_status()
        print("\n📊 数据库状态:")
        print(f"  总节点数: {status.get('total_nodes', 0)}")
        print(f"  疾病数: {status.get('diseases', 0)}")
        print(f"  症状数: {status.get('symptoms', 0)}")
        print(f"  药物数: {status.get('drugs', 0)}")
        print(f"  科室数: {status.get('departments', 0)}")
        print(f"  关系数: {status.get('total_relationships', 0)}")
        
        print("\n✅ 数据库初始化完成！")
        return True
        
    except Exception as e:
        print(f"❌ 初始化失败: {str(e)}")
        logging.error(f"数据库初始化失败: {str(e)}", exc_info=True)
        return False

def show_database_status():
    """显示数据库状态"""
    print("="*60)
    print("📊 Neo4j数据库状态")
    print("="*60)
    
    if not check_neo4j_connection():
        return
    
    try:
        status = neo4j_service.check_database_status()
        
        if 'error' in status:
            print(f"❌ 获取状态失败: {status['error']}")
            return
        
        print(f"数据库URI: {Config.NEO4J_URI}")
        print(f"用户名: {Config.NEO4J_USERNAME}")
        print(f"初始化状态: {'✅ 已初始化' if status.get('total_nodes', 0) > 0 else '❌ 未初始化'}")
        print()
        print("数据统计:")
        print(f"  总节点数: {status.get('total_nodes', 0)}")
        print(f"  疾病数: {status.get('diseases', 0)}")
        print(f"  症状数: {status.get('symptoms', 0)}")
        print(f"  药物数: {status.get('drugs', 0)}")
        print(f"  科室数: {status.get('departments', 0)}")
        print(f"  总关系数: {status.get('total_relationships', 0)}")
        
    except Exception as e:
        print(f"❌ 获取数据库状态失败: {str(e)}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='医疗问答系统 - Neo4j数据库初始化工具')
    parser.add_argument('--init', action='store_true', help='初始化数据库')
    parser.add_argument('--force', action='store_true', help='强制重新初始化（清空现有数据）')
    parser.add_argument('--status', action='store_true', help='显示数据库状态')
    parser.add_argument('--check', action='store_true', help='检查数据库连接')
    
    args = parser.parse_args()
    
    setup_logging()
    
    if args.check:
        check_neo4j_connection()
    elif args.status:
        show_database_status()
    elif args.init:
        success = initialize_database(force=args.force)
        sys.exit(0 if success else 1)
    else:
        # 默认显示状态
        show_database_status()
        print("\n使用帮助:")
        print("  python scripts/init_database.py --init     # 初始化数据库")
        print("  python scripts/init_database.py --force    # 强制重新初始化")
        print("  python scripts/init_database.py --status   # 显示数据库状态")
        print("  python scripts/init_database.py --check    # 检查数据库连接")

if __name__ == '__main__':
    main()